from functools import cache

import numpy as np
import torch
from einops import reduce
from jaxtyping import Float
from torch import Tensor
from .cam_pose_metric import so3_relative_angle_custom as angle_error_mat_batch
from .cam_pose_metric import compare_translation_by_angle as angle_error_vec_batch


def compute_translation_error_batch(t1: torch.Tensor, t2: torch.Tensor) -> torch.Tensor:
    """
    Computes the Euclidean distance between batches of translation vectors.

    Args:
        t1 (torch.Tensor): Batch of translation vectors (B, 3).
        t2 (torch.Tensor): Batch of translation vectors (B, 3).

    Returns:
        torch.Tensor: Batch of Euclidean distances (B,).
    """
    # Compute the difference vector for each item in the batch
    difference = t1 - t2 # (B, 3) - (B, 3) -> (B, 3)

    # Compute the norm of each difference vector along the last dimension (dim=-1)
    # torch.linalg.norm(..., dim=-1) returns a tensor of shape (B,)
    error = torch.linalg.norm(difference, dim=-1)

    return error

@torch.no_grad()
def compute_pose_error_batch(pose_gt: torch.Tensor, pose_pred: torch.Tensor):
    """
    Computes pose errors for a batch of pose matrices using batch-aware helpers.

    Args:
        pose_gt (torch.Tensor): Ground truth pose matrices (B, 4, 4).
        pose_pred (torch.Tensor): Predicted pose matrices (B, 4, 4).

    Returns:
        tuple: (error_t, error_t_scale, error_R)
               error_t (torch.Tensor): Batch of rotation errors for translation vectors (B,).
               error_t_scale (torch.Tensor): Batch of translation scale/distance errors (B,).
               error_R (torch.Tensor): Batch of rotation matrix errors (B,).
    """
    R_gt = pose_gt[:, :3, :3] # Shape (B, 3, 3)
    t_gt = pose_gt[:, :3, 3]  # Shape (B, 3)

    R = pose_pred[:, :3, :3]  # Shape (B, 3, 3)
    t = pose_pred[:, :3, 3]   # Shape (B, 3)

    # Use the batch-aware helper functions
    error_t = angle_error_vec_batch(t, t_gt) # Output shape (B,)

    error_t_scale = compute_translation_error_batch(t, t_gt) # Output shape (B,)

    error_R = angle_error_mat_batch(R, R_gt) # Output shape (B,)

    return error_t, error_t_scale, error_R